﻿#include "RpcCryptoConnection.h"
#include "XyNet/Channel/RpcSocketChannel.h"
#include "XyNet/Common/Cast.h"
#include "XyNet/Common/Exception.h"
#include "XyNet/Common/Logger.h"
#include "XyNet/Endpoint/RpcEndpoint.h"
#include "XyNet/Protocol/cpp/Rpc.service.h"
#include "XyNet/Service/RpcController.h"

namespace XyNet
{
	RpcCryptoConnection::RpcCryptoConnection(RpcSocketChannel* socketChannel)
		: RpcConnection(socketChannel)
	{
		auto config = socketChannel->getEndpoint()->getConfig();
		auto eccType = Ecc::Ecc256;
		auto aesType = Aes::Aes128;

		switch (config->getEncryptionLevel())
		{
		case 1:
			break;
		case 2:
			eccType = Ecc::Ecc384;
			aesType = Aes::Aes192;
			break;
		case 3:
			eccType = Ecc::Ecc512;
			aesType = Aes::Aes256;
			break;
		default:
			Logger::getDefault()->logWarning("无效的加密级别");
		}

		_eccDecryptor = make_unique<Ecc::Decryptor>(Ecc::Ecc256);
		_aesEncryptor = make_unique<Aes>(Aes::Aes128);
		_aesRemoteDecryptor = make_unique<Aes>(Aes::Aes128);
	}

	RpcCryptoConnection::~RpcCryptoConnection() {}

	void RpcCryptoConnection::start()
	{
		// 发送ECC公钥
		Ecc::Encryptor eccEncryptor(*_eccDecryptor.get());
		asyncWrite(eccEncryptor.getPublicKey());

		// 读取对方公钥
		asyncRead(bind(&RpcCryptoConnection::handleReadEccPublicKey,
			XYNET_POINTER_DOWNCAST<RpcCryptoConnection>(shared_from_this()), _1));
	}

	void RpcCryptoConnection::write(const Bytes& data)
	{
		assert(!data.empty());
		if (isConnectComplete())
		{
			// 连接完成后加密器才有效
			Bytes encryptedData = _aesEncryptor->encrypt(data);
			//_aesEncryptor.setRandomIv();
			//encryptedData = _aesEncryptor.getIv() + encryptedData;	// 附加新iv，供远端下次使用
			RpcConnection::write(encryptedData);
		}
		else
		{
			pushCacheWriteData(data);
		}
	}

	void RpcCryptoConnection::read()
	{
		asyncRead(bind(&RpcCryptoConnection::handleReadEncryptedData,
			XYNET_POINTER_DOWNCAST<RpcCryptoConnection>(shared_from_this()), _1));
	}

	void RpcCryptoConnection::handleReadEccPublicKey(const Bytes& data)
	{
		auto socketChannel = getSocketChannel();
		if (!socketChannel)
			return;

		try
		{
			// 生成远端ECC加密器
			_eccRemoteEncryptor = make_unique<Ecc::Encryptor>(data);

			// 加密发送AES密钥
			Bytes ivAndKey = _aesEncryptor->getIv() + _aesEncryptor->getKey();
			asyncWrite(_eccRemoteEncryptor->encrypt(ivAndKey));

			// 读取对方AES KEY
			auto cryptoConnection =
				XYNET_POINTER_DOWNCAST<RpcCryptoConnection>(shared_from_this());
			asyncRead(bind(&RpcCryptoConnection::handleReadAesKey, cryptoConnection, _1));
		}
		catch (const exception& e)
		{
			XYNET_HANDLE_EXCEPTION(e);
			socketChannel->handleSocketChannelException(RpcErrors::InvalidArgument);
		}
	}

	void RpcCryptoConnection::handleReadAesKey(const Bytes& data)
	{
		auto socketChannel = getSocketChannel();
		if (!socketChannel)
			return;

		try
		{
			// 生成远端AES解密器
			Bytes decryptedData = _eccDecryptor->decrypt(data);
			Bytes iv = decryptedData.substr(0, _aesRemoteDecryptor->getIvByteLength());
			Bytes key = decryptedData.substr(_aesRemoteDecryptor->getIvByteLength());
			_aesRemoteDecryptor->setKeyAndIv(key, iv);

			// 提示准备完成
			setConnectComplete();
		}
		catch (const exception& e)
		{
			XYNET_HANDLE_EXCEPTION(e);
			socketChannel->handleSocketChannelException(RpcErrors::InvalidArgument);
		}
	}

	void RpcCryptoConnection::handleReadEncryptedData(const Bytes& data)
	{
		auto socketChannel = getSocketChannel();
		if (!socketChannel)
			return;

		try
		{
			// 解密数据
			Bytes decryptedData = _aesRemoteDecryptor->decrypt(data);

			//Bytes decryptedData =
			//	_aesRemoteDecryptor.decrypt(data.substr(_aesRemoteDecryptor.getIvByteLength()));
			//// 设置下次通讯使用的iv，因为处理数据时可能会收发消息，所以要在这里设置iv
			//_aesRemoteDecryptor.setIv(data.substr(0, _aesRemoteDecryptor.getIvByteLength()));

			// 处理解密的数据
			socketChannel->handleConnectionData(decryptedData);
		}
		catch (const exception& e)
		{
			XYNET_HANDLE_EXCEPTION(e);
			socketChannel->handleSocketChannelException(RpcErrors::InvalidArgument);
		}
	}
}